{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ELzYjgrYY8x5"
      },
      "source": [
        "##### Copyright 2022 The TensorFlow Authors.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jQc0tqMFKBmv"
      },
      "source": [
        "# Image aesthetic scoring with MUSIQ models\n",
        "\n",
        "In this colab, you can try any of the 4 [MUSIQ](https://ai.googleblog.com/2022/10/musiq-assessing-image-aesthetic-and.html) models to score the aesthetic of an image.\n",
        "\n",
        "These models have a specific characteristic that it can handle images in any size or aspect ratio, making it not necessary to do any preprocessing on the image\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2oNh0fGUlHOF"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow_hub as hub\n",
        "\n",
        "import requests\n",
        "from PIL import Image\n",
        "from io import BytesIO\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KypTsmkzoV1y"
      },
      "source": [
        "# Load Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1wO1PVJqlQxr"
      },
      "outputs": [],
      "source": [
        "selected_model = 'koniq' #@param ['spaq', 'koniq', 'paq2piq', 'ava']\n",
        "\n",
        "NAME_TO_HANDLE = {\n",
        "    # Model trained on SPAQ dataset: https://github.com/h4nwei/SPAQ\n",
        "    'spaq': 'https://tfhub.dev/google/musiq/spaq/1',\n",
        "\n",
        "    # Model trained on KonIQ-10K dataset: http://database.mmsp-kn.de/koniq-10k-database.html\n",
        "    'koniq': 'https://tfhub.dev/google/musiq/koniq-10k/1',\n",
        "\n",
        "    # Model trained on PaQ2PiQ dataset: https://github.com/baidut/PaQ-2-PiQ\n",
        "    'paq2piq': 'https://tfhub.dev/google/musiq/paq2piq/1',\n",
        "\n",
        "    # Model trained on AVA dataset: https://ieeexplore.ieee.org/document/6247954\n",
        "    'ava': 'https://tfhub.dev/google/musiq/ava/1',\n",
        "}\n",
        "\n",
        "model_handle = NAME_TO_HANDLE[selected_model]\n",
        "model = hub.load(model_handle)\n",
        "predict_fn = model.signatures['serving_default']\n",
        "\n",
        "print(f'loaded model {selected_model} ({model_handle})')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nuIo92Y-oXSZ"
      },
      "source": [
        "# Util functions\n",
        "\n",
        "Functions to download an image and to show it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_BQJLaepC9YV"
      },
      "outputs": [],
      "source": [
        "def load_image_from_url(img_url):\n",
        "  \"\"\"Returns an image with shape [1, height, width, num_channels].\"\"\"\n",
        "  user_agent = {'User-agent': 'Colab Sample (https://tensorflow.org)'}\n",
        "  response = requests.get(img_url, headers=user_agent)\n",
        "  image_bytes = BytesIO(response.content)\n",
        "  image = Image.open(image_bytes)\n",
        "  return image, response.content\n",
        "\n",
        "def show_image(image, title=''):\n",
        "  image_size = image.size\n",
        "  plt.imshow(image)\n",
        "  plt.axis('on')\n",
        "  plt.title(title)\n",
        "  plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tksog9MdI3li"
      },
      "source": [
        "# Load an image\n",
        "\n",
        "The prediction will run directly on the bytes loaded, no preprocessing is needed"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P0oULq_NnxNg"
      },
      "outputs": [],
      "source": [
        "image_url = 'https://blogger.googleusercontent.com/img/b/R29vZ2xl/AVvXsEgr0DKaAoO6qTrJo3hXP8UM3D4AB8gQeNI22Q2QphBVGgn-5v84tjhH3ZWTlGtlUoPdlcx54dM93Qi04MuN7eBbj9WlT8Qxy6B2Us4kcn_53FH28MnTtGCzMPhjCVGIgXRL8ZEMeO-7iue7sNEGxBtgx2bI-eKDQAondM8Dfjb1FaybFgUQji4UU9-0vQ/s1024/image9.png' #@param {type: 'string'}\n",
        "\n",
        "image, image_bytes = load_image_from_url(image_url)\n",
        "\n",
        "show_image(image)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3-i-cBoEJMn_"
      },
      "source": [
        "# Run Prediction on a Single Image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TuE4kKpvokQZ"
      },
      "outputs": [],
      "source": [
        "prediction = predict_fn(tf.constant(image_bytes))\n",
        "print(\"predicted MOS: \", prediction)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "//learning/deepmind/public/tools/ml_python:ml_notebook",
        "kind": "private"
      },
      "provenance": [
        {
          "file_id": "1i9hws3sIbIAzPFW-XAiRQO9PgVoYDXc3",
          "timestamp": 1667387207919
        }
      ]
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
